#!/usr/local/bin/python

"""
Compiles scripts for running rseg on lustre Submits to lustre 
Input is bamfiles, output will go to dir current dir
"""

import os
import subprocess
import sys
import re

# ----------------------------[ Settings ]--------------------------------------

BAMDIR= '/lustre/sblab/berald01/projects/20111212_chipseq_formylseq_raiberea/bam_clean'             ## Dir with bamfiles
CTRLFILE= 'bham385.bwa.mm9.bam'                  ## The control file to use. Assumed to be in *BAMDIR*. Note that it must be a *BED* file
SKIPBAMS= ['bham385.bwa.mm9.bam', 'j1_input_SRR185880.bwa.mm9.mapq15.bam'] ## List of bamfiles in BAMDIR to skip (e.g. the control file)

OUTDIR= '/lustre/sblab/berald01/projects/20111212_chipseq_formylseq_raiberea/rseg/mode2_bin-100bp.2' ## Where output, including sh scripts and log files will be

# ------------------------------------------------------------------------------

p= subprocess.Popen('mkdir -p %s' %(OUTDIR), shell= True)
p.wait()

## p= subprocess.Popen('mkdir -p rseg') ## Where output files will be
bamfiles= sorted(os.listdir(BAMDIR))
bamfiles= [x for x in bamfiles if re.search('\.bam$', x)]
ctrlfile_path= os.path.join(BAMDIR, CTRLFILE)

n= 1
for b in bamfiles:
    if b in SKIPBAMS:
        continue
    bam= os.path.join(BAMDIR, b)
    shfilename= os.path.join(OUTDIR, re.sub('\.bam$', '', b) + '.py')
    shfile= open(shfilename, 'w')
    cmd= """#!/usr/local/bin/python

import subprocess
import os
import shutil

## Reads in control
cmd= "samtools view -c %(ctrlfile_path)s"
print('\\n\\nExecuting:\\n' + cmd)
p= subprocess.Popen(cmd, shell= True, stdout= subprocess.PIPE)
p.wait()
ctrl_count= float(p.stdout.read().strip())

## Reads in pulldown
cmd= "samtools view -c %(bam)s"
print('\\n\\nExecuting:\\n' + cmd)
p= subprocess.Popen(cmd, shell= True, stdout= subprocess.PIPE)
p.wait()
chip_count= float(p.stdout.read().strip())

print('Reads in control: ' + str(ctrl_count) + ' Reads in pulldown: ' + str(chip_count))

if chip_count >= ctrl_count:
    shutil.copyfile('%(ctrlfile_path)s', '%(ctrlsampled)s')
else:
    ratio= chip_count / ctrl_count
    cmd= 'java -jar /home/berald01/bin/DownsampleSam.jar INPUT=%(ctrlfile_path)s OUTPUT=%(ctrlsampled)s PROBABILITY=' + str(ratio) + ' VALIDATION_STRINGENCY=SILENT RANDOM_SEED=1234'
    print('\\n\\nExecuting:\\n' + cmd)
    p= subprocess.Popen(cmd, shell= True)
    p.wait()

cmd= 'bamToBed -i %(ctrlsampled)s > %(ctrlbed)s'
print('\\n\\nExecuting:\\n' + cmd)
p= subprocess.Popen(cmd, shell= True)    
p.wait()

cmd= 'bamToBed -i %(bam)s > %(chipbed)s'
print('\\n\\nExecuting:\\n' + cmd)
p= subprocess.Popen(cmd, shell= True)
p.wait()

cmd= '/home/berald01/applications/rseg/rseg-bin/bin/rseg-diff -posterior-cutoff 0.9 -v -o %(OUTDIR)s -d /home/berald01/applications/rseg/auxiliary_files/deadzones-k40-mm9.bed -c /home/berald01/applications/rseg/auxiliary_files/mouse-mm9-size.bed -mode 2 -b 100 %(chipbed)s %(ctrlbed)s'
print('\\n\\nExecuting:\\n' + cmd)
p= subprocess.Popen(cmd, shell= True)
p.wait()

os.remove('%(ctrlsampled)s')
os.remove('%(ctrlbed)s')
os.remove('%(chipbed)s')

                """ %{'bam':bam,
                      'ctrlfile_path':ctrlfile_path,
                      'ctrlsampled': ctrlfile_path + '-job-' + str(n) + '.bam',
                      'ctrlbed': ctrlfile_path + 'job-' + str(n) + '.bed',
                      'chipbed': bam + '.bed',
                      'OUTDIR':OUTDIR}
    shfile.write(cmd)
    shfile.close()
    os.chmod(shfilename, 0744)
    jobid= b + str(n)
    cmd= """bsub -R "rusage[mem=8192]" -o %(shfilename)s.log -J %(jobid)s < %(shfilename)s""" %{'jobid':jobid, 'shfilename':shfilename}
    print(cmd)
    p= subprocess.Popen(cmd, shell= True)
    n += 1
sys.exit()

